/**
 * @file VonMisesPlasticity.cpp
 * @brief Implements Von Mises plasticity model for shell elements
 * 
 */

#include <vector>
#include <cmath>
#include <iostream>
#include <string>
#include <stdexcept>

// Constants
constexpr double THIRD = 1.0 / 3.0;
constexpr double CONVERGENCE_TOL = 0.001;
constexpr int MAX_ITERATIONS = 50;

/**
 * @brief Implements Von Mises plasticity model for shell elements
 * 
 * @param numPlasticElements Number of plastic elements
 * @param plasticElementIndices Indices of plastic elements
 * @param yieldStress Yield stress values
 * @param a0 Material parameter A0
 * @param a1 Material parameter A1
 * @param a2 Material parameter A2
 * @param shearModulus Shear modulus values
 * @param plasticStrainIncrement Plastic strain increment
 * @param equivalentStress Equivalent stress (Von Mises)
 * @param stress Stress tensor (updated in-place)
 * @param effectivePlasticStrain Effective plastic strain (updated)
 * @param strainIncrement Strain increment tensor
 * @param maxElements Maximum number of elements in group
 */
void VonMisesPlasticity(
    int& numPlasticElements,
    const std::vector<int>& plasticElementIndices,
    const std::vector<double>& yieldStress,
    const std::vector<double>& a0,
    const std::vector<double>& a1,
    const std::vector<double>& a2,
    const std::vector<double>& shearModulus,
    std::vector<double>& plasticStrainIncrement,
    std::vector<double>& equivalentStress,
    std::vector<std::array<double, 5>>& stress,
    std::vector<double>& effectivePlasticStrain,
    std::vector<std::array<double, 6>>& strainIncrement,
    int maxElements)
{
    // Temporary storage for stress components during iteration
    std::vector<double> sig1(maxElements), sig2(maxElements), sig3(maxElements);
    std::vector<double> sig4(maxElements), sig5(maxElements);
    std::vector<double> a3(maxElements), cons1(maxElements), cons2(maxElements);
    std::vector<double> cons3(maxElements), depsi(maxElements);
    
    // Active elements during iteration
    std::vector<int> activeElementIndices(maxElements);
    int numActiveElements = numPlasticElements;
    
    // Initialize with radial return mapping
    for (int j = 0; j < numPlasticElements; ++j) {
        int i = plasticElementIndices[j];
        double yinv = 1.0 / yieldStress[i];
        
        // Precompute constants
        a3[i] = -0.5 * (a0[i] + 1.0) * yinv;
        cons1[i] = yinv * (a1[i] - 0.5 * a2[i]);
        cons2[i] = yinv * (a2[i] - 0.5 * a1[i]);
        cons3[i] = yinv * 3.0 * shearModulus[i];
        
        // Initial plastic strain increment estimate
        plasticStrainIncrement[i] = THIRD * (equivalentStress[i] - yieldStress[i]) / shearModulus[i];
        depsi[i] = plasticStrainIncrement[i];
        
        // First iteration
        double fac1 = plasticStrainIncrement[i] * cons1[i] + 1.0;
        double fac2 = plasticStrainIncrement[i] * cons2[i];
        double fac = fac1 * fac1 - fac2 * fac2;
        fac1 /= fac;
        fac2 = -fac2 / fac;
        
        // Update stress components
        sig1[i] = fac1 * stress[i][0] + fac2 * stress[i][1];
        sig2[i] = fac2 * stress[i][0] + fac1 * stress[i][1];
        
        // Compute equivalent stress
        double s2v = sig1[i] * (sig1[i] - sig2[i]) + sig2[i] * sig2[i];
        
        double fac3 = 1.0 / (plasticStrainIncrement[i] * cons3[i] + 1.0);
        sig3[i] = stress[i][2] * fac3;
        sig4[i] = stress[i][3] * fac3;
        sig5[i] = stress[i][4] * fac3;
        
        s2v = std::sqrt(s2v + 3.0 * (sig3[i]*sig3[i] + sig4[i]*sig4[i] + sig5[i]*sig5[i]));
        
        // Check convergence
        double yms2 = yieldStress[i] - s2v;
        double conv = std::abs(yms2) * yinv;
        
        if (conv > CONVERGENCE_TOL && plasticStrainIncrement[i] >= 0.0) {
            activeElementIndices[numActiveElements] = i;
            depsi[i] = yms2 * depsi[i] / (s2v - equivalentStress[i]);
            plasticStrainIncrement[i] += depsi[i];
            equivalentStress[i] = s2v;
            numActiveElements++;
        }
    }
    
    // Return if all elements converged
    if (numActiveElements == 0) return;
    
    // Iterative plastic correction
    for (int iter = 0; iter < MAX_ITERATIONS; ++iter) {
        int newActiveElements = 0;
        
        for (int j = 0; j < numActiveElements; ++j) {
            int i = activeElementIndices[j];
            
            // Compute factors
            double fac1 = plasticStrainIncrement[i] * cons1[i] + 1.0;
            double fac2 = plasticStrainIncrement[i] * cons2[i];
            double fac = fac1 * fac1 - fac2 * fac2;
            fac1 /= fac;
            fac2 = -fac2 / fac;
            
            // Update stress components
            sig1[i] = fac1 * stress[i][0] + fac2 * stress[i][1];
            sig2[i] = fac2 * stress[i][0] + fac1 * stress[i][1];
            
            // Compute equivalent stress
            double s2v = sig1[i] * (sig1[i] - sig2[i]) + sig2[i] * sig2[i];
            
            double fac3 = 1.0 / (plasticStrainIncrement[i] * cons3[i] + 1.0);
            sig3[i] = stress[i][2] * fac3;
            sig4[i] = stress[i][3] * fac3;
            sig5[i] = stress[i][4] * fac3;
            
            s2v = std::sqrt(s2v + 3.0 * (sig3[i]*sig3[i] + sig4[i]*sig4[i] + sig5[i]*sig5[i]));
            
            // Check convergence
            double yms2 = yieldStress[i] - s2v;
            double conv = std::abs(yms2) / yieldStress[i];
            
            if (conv > CONVERGENCE_TOL && plasticStrainIncrement[i] >= 0.0) {
                activeElementIndices[newActiveElements] = i;
                depsi[i] = yms2 * depsi[i] / (s2v - equivalentStress[i]);
                plasticStrainIncrement[i] += depsi[i];
                equivalentStress[i] = s2v;
                newActiveElements++;
            }
        }
        
        numActiveElements = newActiveElements;
        if (numActiveElements == 0) break;
    }
    
    // Warning if not all elements converged
    if (numActiveElements > 0) {
        std::cerr << "*** WARNING *** NO CONVERGENCE ACHIEVED AFTER " 
                  << MAX_ITERATIONS << " ITERATIONS FOR " 
                  << numActiveElements << " ELEMENTS\n";
    }
    
    // Final update of stresses and strains
    for (int j = 0; j < numPlasticElements; ++j) {
        int i = plasticElementIndices[j];
        
        // Update stress components
        stress[i][0] = sig1[i];
        stress[i][1] = sig2[i];
        stress[i][2] = sig3[i];
        stress[i][3] = sig4[i];
        stress[i][4] = sig5[i];
        
        // Update plastic strain
        effectivePlasticStrain[i] += plasticStrainIncrement[i];
        
        // Update out-of-plane strain
        strainIncrement[i][5] += a3[i] * plasticStrainIncrement[i] * 
                                (stress[i][0] + stress[i][1]);
        
        // Check for negative plastic strain
        if (plasticStrainIncrement[i] < 0.0) {
            throw std::runtime_error("NEGATIVE PLASTIC STRAIN FOR USR-MAT SHELL ELEMENT");
        }
    }
}

/**
 * @brief Example usage of the Von Mises plasticity model
 */
void ExampleUsage() {
    const int maxElements = 100;
    int numPlasticElements = 10;
    
    std::vector<int> plasticElementIndices(maxElements);
    std::vector<double> yieldStress(maxElements, 250.0);
    std::vector<double> a0(maxElements, 1.0), a1(maxElements, 1.0), a2(maxElements, 1.0);
    std::vector<double> shearModulus(maxElements, 80000.0);
    std::vector<double> plasticStrainIncrement(maxElements);
    std::vector<double> equivalentStress(maxElements, 300.0); // Above yield
    std::vector<std::array<double, 5>> stress(maxElements);
    std::vector<double> effectivePlasticStrain(maxElements);
    std::vector<std::array<double, 6>> strainIncrement(maxElements);
    
    // Initialize some example data
    for (int i = 0; i < maxElements; ++i) {
        plasticElementIndices[i] = i;
        stress[i] = {300.0, 250.0, 100.0, 50.0, 50.0}; // Sample stress state
    }
    
    try {
        VonMisesPlasticity(
            numPlasticElements,
            plasticElementIndices,
            yieldStress,
            a0,
            a1,
            a2,
            shearModulus,
            plasticStrainIncrement,
            equivalentStress,
            stress,
            effectivePlasticStrain,
            strainIncrement,
            maxElements);
        
        std::cout << "Plasticity update completed. Number of plastic elements: " 
                  << numPlasticElements << std::endl;
    } catch (const std::exception& e) {
        std::cerr << "Error in plasticity calculation: " << e.what() << std::endl;
    }
}